#!/usr/bin/python
# Copyright (c) 2011 Josef Robert Novak
#
# You may copy and modify this freely under the same terms as
# Sphinx-III
#
#  This module converts an arbitrary Sphinx format acoustic model to 
#  to format suitable for use with the TCubed WFST decoder.  
#  The text format is also very similar to the AT&T standard.
#
#  The Sphinx AM code comes directly from the sphinxtrain python
#  libraries and has not been significantly modified.  
#
#  t3amConvert.py (this file) is the only real original contribution.
#  It reads in the binary Sphinx acoustic models and transforms them
#  into a standard intermediate format.  TCubed uses this same format
#  for HTK models as well.

__author__ = "Josef Robert Novak <novakj@gavo.t.u-tokyo.ac.jp>"
__version__ = "$Revision $"

from t3model import T3Model
import os



class sphinx2T3():
    """
    This class will read in an arbitrary Sphinx format Acoustic Model and convert
    it to the TCubed text-format. It relies on the Sphinx python libraries to read in the 
    binary acoustic models.
    """
    def __init__( self, modelroot="./", outputroot="./t3py", hmmhmm="hmm.hmm", carcs="clevel.arcs", 
                  hmmdur="hmm.dur", hmmst="hmm.st", hmmmean="hmm.mean", hmmsyms="hmm.syms",
                  hmmvar="hmm.var", hlevelarcs="hlevel.arcs", hmmpdf="hmm.pdf" ):
        #Setup the output files
        self.modelroot  = modelroot
        self.outputroot = outputroot
        if not os.path.exists(self.outputroot): os.mkdir(self.outputroot)
        self.hmmhmm     = os.path.join( outputroot, hmmhmm )
        self.carcs      = os.path.join( outputroot, carcs )
        self.hmmdur     = os.path.join( outputroot, hmmdur )
        self.hmmst      = os.path.join( outputroot, hmmst )
        self.hmmmean    = os.path.join( outputroot, hmmmean )
        self.hmmvar     = os.path.join( outputroot, hmmvar )
        self.hlevelarcs = os.path.join( outputroot, hlevelarcs )
        self.hmmpdf     = os.path.join( outputroot, hmmpdf )
        self.hmmsyms    = os.path.join( outputroot, hmmsyms )
        #Load the binary Sphinx AM
        self.sphinx_am  = T3Model( self.modelroot )
        self._compute_durations( )
        self.state2durations = {}
        
    def build_all( self ):
        """Convenience function to generate all the models."""
        self.generateHMMsyms( self.hmmsyms )
        self.generateHMMandCLevelArcs( self.carcs, self.hmmhmm )
        self.generateEverythingElse( self.hmmdur, self.hmmst, self.hmmmean, self.hmmvar, self.hlevelarcs, self.hmmpdf )
        return

    def generateHMMsyms( self, hmmsyms ):
        """
        Generate the HMM syms file from the mdef.  
        This just formats the phones/cd_phones and writes out in the order
        specified by the mdef file.  The epsilon symbol '-' is given priority.
        """
        hmmsyms_fp = open( hmmsyms, "w" )
        hmmsyms_fp.write("- 0\n")
        for i,ci_cd in enumerate(self.sphinx_am.mdef.allfields):
            phone = self._format_phone( ci_cd[0:4] )
            hmmsyms_fp.write("%s %d\n" % ( phone, i+1 ))
        hmmsyms_fp.close()
        return

    def _format_phone( self, ci_cd ):
        """Format the mdef phones.  This convention will need to be followed in all top level transducers."""
        #monophone case
        if ci_cd[1] == "-":
            formatted = ci_cd[0]
        #context-dependent triphone case
        else:
            formatted = ci_cd[1] + '-' + ci_cd[0] + '+' + ci_cd[2] + '_' + ci_cd[3]
        return formatted

    def generateEverythingElse( self, hmmdur, hmmst, hmmmean, hmmvar, hlevelarcs, hmmpdf ):
        """Generate all remaining files as needed."""

        #Open the output files
        hmmdur_fp     = open( hmmdur, "w" )
        hmmst_fp      = open( hmmst, "w" )
        hmmmean_fp    = open( hmmmean, "w" )
        hmmvar_fp     = open( hmmvar, "w" )
        hlevelarcs_fp = open( hlevelarcs, "w" )
        hmmpdf_fp     = open( hmmpdf, "w" )
        
        #Now to the magic!
        hlevelarcs_fp.write("%d 2\n" % (len(self.state2durations)) )
        for i in xrange( len(self.state2durations) ):
            f = self.state2durations[i]
            durMean = 1.0 / (1.0 - f[0])
            durVar = durMean * durMean * 0.5
            hmmdur_fp.write("%d\t-1\t%0.7f\t%0.7f\n" % ( i+1, durMean, durVar ))
           
            hlevelarcs_fp.write("%d 0 1 x\n" % ( i+1 ))
            hlevelarcs_fp.write("%d %0.7f %0.7f\n" % ( i+1, f[0], f[1] ))

            hmmst_fp.write("%d\t-1" % ( i+1 ))
            gaussianPerState = len(self.sphinx_am.mixw[0][0])
            
            for j in xrange(gaussianPerState):
                componentID = i * gaussianPerState + j + 1
                hmmst_fp.write("\t%d\t%0.7f" % (componentID, self.sphinx_am.mixw[i][0][j]))
                hmmmean_fp.write("%d\t-1" % (componentID) )
                hmmvar_fp.write("%d\t-1"  % (componentID) )
                
                dimensions = len(self.sphinx_am.mean[0][0][0])
                for k in xrange(dimensions):
                    hmmmean_fp.write("\t%0.7f" % ( self.sphinx_am.mean[i][0][j][k] ) )
                    #Unwind the inversion done in the s3model.py code: this is done in the decoder
                    hmmvar_fp.write("\t%0.7f"  % ( (1./self.sphinx_am.var[i][0][j][k])/2.) )

                hmmmean_fp.write("\n")
                hmmvar_fp.write("\n")
                hmmpdf_fp.write("%d\t%d\t%d\t%d\n" % (componentID, componentID, componentID, 0) )
            hmmst_fp.write("\n")
        #Close up shop
        hmmdur_fp.close()
        hmmst_fp.close()
        hmmmean_fp.close()
        hmmvar_fp.close()
        hlevelarcs_fp.close()
        hmmpdf_fp.close()

    def generateHMMandCLevelArcs( self, carcs, hmmhmm ):
        """Generate hmm.hmm, and the clevel arcs."""

        carcs_fp  = open( carcs, "w" )
        hmmhmm_fp = open( hmmhmm, "w" )
        carcs_fp.write("%d 2\n" % ( self.sphinx_am.mdef.n_phone ))
        
        for i, fields in enumerate(self.sphinx_am.mdef.allfields):
            carcs_fp.write("%d 0 3 x\n" % ( i+1 ))
            hmmhmm_fp.write("%d\t1" % ( i+1 ))
            
            for j,stateID in enumerate(fields[6:-1]):
                carcs_fp.write("%d\t%0.10f\t%0.7f\n" % (int(stateID)+1, self.durations[int(fields[5])][j][0], self.durations[int(fields[5])][j][1]))
                hmmhmm_fp.write("\t%d" % ( int(stateID)+1 ))
                
                if not self.state2durations.has_key( int(stateID) ):
                    self.state2durations[int(stateID)] = (self.durations[int(fields[5])][j][0], self.durations[int(fields[5])][j][1])
                    
            hmmhmm_fp.write("\n")
        hmmhmm_fp.close()
        carcs_fp.close()
        return

    def _compute_durations( self ):
        """Compute the durations from the transition matrices."""
        self.durations = [] #Unfortunately we can't initialize arrays nicely in python
        for i,tmat in enumerate(self.sphinx_am.tmat):
            self.durations.append([])
            for j,row in enumerate(tmat):
                self.durations[i].append([])
                sum = 0.0
                for k,state in enumerate(row):
                    sum += state
                for k,state in enumerate(row):
                    if j==k:
                        self.durations[i][j].append( state/sum )
                    elif j==k-1:
                        self.durations[i][j].append( state/sum )
        return


if __name__=="__main__":
    import sys
    sphinxModel = sphinx2T3( sys.argv[1] )
    sphinxModel.build_all( )
